from ship_envs.core.action import Action
from ship_envs.core.fairway import Fairway
from ship_envs.core.ship import Ship
from ship_envs.core.ship_forward import ShipForward
from ship_envs.environment.single_ship_env import SingleShipEnv


class ImazuProblemEnv(SingleShipEnv):
    def __init__(self, case):
        super().__init__()
        self.agent_ship_pos, self.agent_ship_angle, self.agent_ship_target_pos, self.agent_ship_speed, \
        self.other_ship_pos, self.other_ship_angle, self.other_ship_speed = case
        for _ in self.other_ship_pos:
            self.world.add_ship(ShipForward())
        self.reward_threshold = 5

    def reset(self):
        # 重置船位置及目标点位置
        ship = self.get_ship()
        ship.cur_speed = self.agent_ship_speed[0]
        ship.set_pos(self.agent_ship_pos[0])
        ship.set_angle(self.agent_ship_angle[0])
        ship.start_angle = self.agent_ship_angle[0]
        # 干扰船
        for i in range(len(self.other_ship_pos)):
            other_ship = self.world.get_other_ship(i)
            other_ship.cur_speed = self.other_ship_speed[i]
            pos = self.other_ship_pos[i]
            angle = self.other_ship_angle[i]
            other_ship.set_pos(pos)
            other_ship.set_angle(angle)

        # 重置步数
        self.step_count = 0
        # 刷新obs
        self.world.refresh()
        # 单智能体，obs为整个world的obs  全观测
        return self.observation_callback(self.world, ship)

    def reward_callback(self, world, ship):
        if ship.out_range():
            if ship.dif_with_start_angle_abs() == 0:
                return 10
        if world.has_collision():
            return -10
        reward = -ship.dif_with_start_angle_abs() / 90
        # 多左转做一定惩罚，使得智能体优先右转
        if ship.cur_action == Action.TURNLEFT.value:
            reward -= 0.005
        return reward
